#include "refine/smooth3.h"
#include "refine/window_pass.h"
#include <limits>
#include <iostream>
#include <set>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g3/smatrix.h>
//#include <glib/g3/lmatrix.h> //temporary

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth3 {

const params default_params = {2, 2, 1, 100/*, 100*/, false, false};

typedef grey_mask::mask::elem_type grey_mask_value;
const grey_mask_value grey_mask_max_value = std::numeric_limits<grey_mask_value>::max();	//255
const grey_mask_value grey_mask_min_value = std::numeric_limits<grey_mask_value>::min();	//0
const grey_mask_value grey_mask_avg_value = (grey_mask_max_value + grey_mask_min_value) / 2;//127

typedef unsigned uint;
typedef g3::basic_vector<uint>		vectoru;
typedef g3::basic_smatrix<uint>		smatrixu;
typedef g3::basic_vector<double>	vector;
typedef g3::basic_smatrix<double>	smatrix;
typedef g3::basic_vector<float>		vectorf;

template <typename T>
inline T sqr(T t)
	{ return t * t; }

inline vectoru v(byte_color const & c)
	{ return vectoru(c.b, c.g, c.r); }

/// reconstructs linear relation between alpha of the central pixel and 
/// alphas of neighbour pixels with coefficient depending from image colors in the form
/// alpha_0 = sum_i c(I_i) alpha_i, 
/// where c(I_i) = p + q * I_i
class linear_alpha
{
public:
	linear_alpha()
		{ remove_all(); }
	uint count() const
		{ return n; }
	void remove_all() {
		n = 0;
		sum_I = vectoru::nill;
		sum_IIT.zero();
	}
	void expansion(vector const & I_0, float & p, vectorf & q) const {
		assert (n > 0);
		//double precision here is used for better matrix inversion
		vector qq = sle0(
			double(n) * smatrix(sum_IIT) - tsquare(vector(sum_I)),
			double(n) * I_0 - vector(sum_I));
		q = qq;
		p = (1 - scalar_product(qq, vector(sum_I))) / n;
	}
	void add_sample(vectoru const & I) {
		++n;
		sum_I += I;
		sum_IIT += tsquare(I);
	}
	void remove_sample(vectoru const & I) {
		--n;
		sum_I -= I;
		sum_IIT -= tsquare(I);
	}
protected:
	uint n;
	vectoru sum_I;
	smatrixu sum_IIT; ///< sum of (I * I^T)
};

class linear_alpha_classifier : public linear_alpha
{
public:
	linear_alpha_classifier(/*float max_condition_number_*/) 
		//: max_condition_number(max_condition_number_)
		{ remove_all(); }
	uint count() const
		{ return linear_alpha::count(); }
	void remove_all() {
		linear_alpha::remove_all();
		sum_alpha = sum_sqr_alpha = 0;
		sum_alpha_I = vectoru::nill;
	}
	void add_sample(vectoru const & I, uint alpha) {
		linear_alpha::add_sample(I);
		sum_alpha += alpha;
		sum_sqr_alpha += alpha * alpha;
		sum_alpha_I += alpha * I;
	}
	void remove_sample(vectoru const & I, uint alpha) {
		linear_alpha::remove_sample(I);
		sum_alpha -= alpha;
		sum_sqr_alpha -= alpha * alpha;
		sum_alpha_I -= alpha * I;
	}
	///checks max_condition_number and returns true in the case of fail
	/*bool bad_condition_number() const {
		if (n == 0)
			return true;
		double a = sum_sqr_alpha;
		double b = grey_mask_max_value * sum_alpha - sum_sqr_alpha;
		double c = grey_mask_max_value * (grey_mask_max_value * n - 2 * sum_alpha) + sum_sqr_alpha;
		double D = sqr(a - c) + 4 * b * b;
		return sqr(max_condition_number + 1) * D > sqr((max_condition_number - 1) * (a + c));
	}*/

	/// guesses alpha for the given color;
	/// does not change (g) in the case of fail
	bool classify(vector const & I, grey_mask_value & g/*, bool test_condition_number = true*/) const {
		assert (n > 0);
		/*if (test_condition_number)
			if (bad_condition_number())
				return false;*/

		float p;
		vectorf q;
		expansion(I, p, q);
		classify_using_known_expansion(p, q, g);
		return true;
	}
	void classify_using_known_expansion(
		float p, vectorf const & q, grey_mask_value & g) const 
	{
		float a = p * float(sum_alpha) + scalar_product(q, vectorf(sum_alpha_I)) + 0.5f; //0.5 for proper rounding

		if (a <= grey_mask_min_value)
			g = grey_mask_min_value;
		else if (a >= grey_mask_max_value)
			g = grey_mask_max_value;
		else 
			g = grey_mask_value(a);
	}

	//alpha of any color I can be determined as (a0 + grad_a*I)
	void expansion2(float & a0, vectorf & grad_a) const {
		grad_a = sle0(
			double(n) * smatrix(sum_IIT) - tsquare(vector(sum_I)),
			double(n) * vector(sum_alpha_I) - double(sum_alpha) * vector(sum_I));
		a0 = (sum_alpha - scalar_product(vectorf(sum_I), grad_a)) / n;
	}
	static void classify_using_known_expansion2(
		float a0, vectorf const & grad_a, vectorf const & I, grey_mask_value & g)
	{
		float a = a0 + scalar_product(I, grad_a) + 0.5; //0.5 for proper rounding

		if (a <= grey_mask_min_value)
			g = grey_mask_min_value;
		else if (a >= grey_mask_max_value)
			g = grey_mask_max_value;
		else 
			g = grey_mask_value(a);
	}

private:
	uint sum_alpha, sum_sqr_alpha;
	vectoru sum_alpha_I;
	//float max_condition_number;
};

//template <typename Raster, typename Matte>
class smoother
{
public:
	smoother(byte_image const & src_, 
		grey_mask::mask const & matte_, grey_mask::mask & next_matte_,
		img::bitmap & modifiable_/*, float max_condition_number_*/)
		: src(src_), matte(matte_), next_matte(next_matte_), modifiable(modifiable_)
		  //,classifier(max_condition_number_)
		{ clear(); }
	void clear() {
		classifier.remove_all();
		checks = achanges = 0;
	}
	size_extent dim() const
		{ return src.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p) 
		{ classifier.add_sample(v(src[p]), matte[p]); }
	void remove(size_point const & p) 
		{ classifier.remove_sample(v(src[p]), matte[p]); }
	void remove_all() 
		{ classifier.remove_all(); }
	void process(size_point const & p, size_rect const & wnd) {
		++checks;
		grey_mask_value g;
		if (!classifier.classify(v(src[p]), g))
			modifiable[p] = false;
		else {
			bool not_greater_exists = false;
			bool not_less_exists = false;
			size_point i;
			for (i.y = wnd.y1; i.y < wnd.y2; ++i.y)
				for (i.x = wnd.x1; i.x < wnd.x2; ++i.x) {
					if (!not_greater_exists && matte[i] <= g)
						not_greater_exists = true;
					if (!not_less_exists && matte[i] >= g)
						not_less_exists = true;
					if (not_greater_exists && not_less_exists)
						break;
				}

			if (not_greater_exists && not_less_exists && matte[p] != g) {
				++achanges;
				next_matte[p] = g;
			}
			else
				modifiable[p] = false;
		}
	}
private:
	byte_image const & src;
	grey_mask::mask const & matte;
	grey_mask::mask & next_matte;
	img::bitmap & modifiable;
	linear_alpha_classifier classifier;
	uint checks, achanges;
};

class ultra_clear_smoother
{
public:
	ultra_clear_smoother(byte_image const & src_, 
		grey_mask::mask const & matte_, grey_mask::mask & next_matte_,
		img::bitmap & modifiable_/*, float max_condition_number_*/)
		: src(src_), matte(matte_), next_matte(next_matte_), modifiable(modifiable_)
		  //,classifier(max_condition_number_)
		{ clear(); }
	void clear() {
		classifier.remove_all();
		checks = achanges = 0;
	}
	size_extent dim() const
		{ return src.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p) {
		grey_mask_value g = matte[p];
		byte_color c = src[p];
		classifier.add_sample(v(c), g);
		if (g == grey_mask_min_value)
			back_colors.insert(c);
		else if (g == grey_mask_max_value)
			fore_colors.insert(c);
	}
	void remove(size_point const & p) {
		grey_mask_value g = matte[p];
		byte_color c = src[p];
		classifier.remove_sample(v(c), g);
		if (g == grey_mask_min_value) {
			//simple call to erase method would erase all elements with this color
			colors_set::iterator i = back_colors.find(c);
			assert (i != back_colors.end());
			back_colors.erase(i);
		}
		else if (g == grey_mask_max_value) {
			colors_set::iterator i = fore_colors.find(c);
			assert (i != fore_colors.end());
			fore_colors.erase(i);
		}
	}
	void remove_all() {
		classifier.remove_all(); 
		back_colors.clear();
		fore_colors.clear();
	}
	void process(size_point const & p, size_rect const &) {
		++checks;

		typedef std::vector<byte_color> colors_vector;
		colors_vector back_colors(this->back_colors.begin(), this->back_colors.end());
		colors_vector fore_colors(this->fore_colors.begin(), this->fore_colors.end());
		linear_alpha_classifier classifier(this->classifier);
		bool classified = false;
		grey_mask_value g;
		vectorf I = v(src[p]);
		//while (!classifier.bad_condition_number()) {
		for (;;) {
			float a0;
			vectorf grad_a;
			classifier.expansion2(a0, grad_a);
			classifier.classify_using_known_expansion2(a0, grad_a, I, g);
			classified = true;

			/// leave only those foreground and background colors,
			/// for which classification made mistakes
			colors_vector new_back_colors;
			new_back_colors.reserve(back_colors.size());
			for (uint i = 0; i < back_colors.size(); ++i) {
				if (a0 + scalar_product(vectorf(v(back_colors[i])), grad_a) > grey_mask_min_value)
					new_back_colors.push_back(back_colors[i]);
				else
					classifier.remove_sample(v(back_colors[i]), grey_mask_min_value);
			}

			colors_vector new_fore_colors;
			new_fore_colors.reserve(fore_colors.size());
			for (uint i = 0; i < fore_colors.size(); ++i) {
				if (a0 + scalar_product(vectorf(v(fore_colors[i])), grad_a) < grey_mask_max_value)
					new_fore_colors.push_back(fore_colors[i]);
				else
					classifier.remove_sample(v(fore_colors[i]), grey_mask_max_value);
			}

			if (new_back_colors.size() == back_colors.size() && new_fore_colors.size() == fore_colors.size())
				break;
			back_colors.swap(new_back_colors);
			fore_colors.swap(new_fore_colors);
		}

		if (!classified || matte[p] == g)
			modifiable[p] = false;
		else {
			++achanges;
			next_matte[p] = g;
		}
	}
private:
	byte_image const & src;
	grey_mask::mask const & matte;
	grey_mask::mask & next_matte;
	img::bitmap & modifiable;
	linear_alpha_classifier classifier;
	uint checks, achanges;
private:
	typedef std::multiset<byte_color> colors_set;
	colors_set back_colors, fore_colors;
};

void perform_core(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const img::bitmap & strict_fore, 
	const img::bitmap & strict_back,
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Tupoy Volosok(R) matting..." << endl;

	size_extent dim = src.dim();
	assert (dim == matte.dim());
	assert (dim == modifiable.dim());
	int_rect frame(0, 0, dim.sx, dim.sy);

	grey_mask::mask next_matte = matte;

	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		img::combine(modifiable, strict_fore, img::bit_and_not());
		img::combine(modifiable, strict_back, img::bit_and_not());

		if (p.ultra_clear) {
			ultra_clear_smoother sm(src, matte, next_matte, modifiable/*, p.max_condition_number*/);
			window_pass(sm, p.whs);
			if (!p.silent)
				clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;
			if (sm.change_count() == 0)
				break;
		}
		else {
			smoother sm(src, matte, next_matte, modifiable/*, p.max_condition_number*/);
			window_pass(sm, p.whs);
			if (!p.silent)
				clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;
			if (sm.change_count() == 0)
				break;
		}

		matte = next_matte;

		img::dilate8(modifiable, p.whs);
	}
}


void perform(
	const byte_image & src, 
	const bit_mask::mask & m, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = src.dim();
	assert (dim == m.dim());

	matte.resize(dim);
	for (grey_mask::mask::iterator i = matte.begin(); i != matte.end(); ++i)
		*i = m[i] ? grey_mask_max_value : grey_mask_min_value;

	img::bitmap strict_fore, strict_back;
	img::erode4(strict_fore, m, p.fore_depth);
	img::dilate4(strict_back, m, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable = m; 
	{
		uint w = 1 + p.whs;// * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, w);
		img::bitmap tmp = m;
		tmp.inverse();
		img::dilate4(tmp, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = matte.dim();

	img::bitmap not_min(dim), not_max(dim), strict_fore(dim), strict_back(dim);
	for (grey_mask::mask::const_iterator i = matte.begin(); i != matte.end(); ++i) {
		not_min[i] = *i != grey_mask_min_value;
		not_max[i] = *i != grey_mask_max_value;
		strict_fore[i] = *i > grey_mask_avg_value;
	}

	strict_back = strict_fore;
	img::erode4(strict_fore, p.fore_depth);
	img::dilate4(strict_back, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs;// * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	argb_byte_image & src,
	const params & p)
{
	//!!! temporary solution
	byte_image src2 = src; 
	grey_mask::mask matte = img::projection(src, &argb_byte_color::a);

	perform(src2, matte, p);

	img::projection(src, &argb_byte_color::a) = matte;
}

} //namespace smooth3

} //namespace refine
